from .resnet import *
import torch.nn as nn
import torch
from .globalNet import globalNet
from .refineNet import refineNet
from .simple_pose import get_pose_net
from .Unet import UNet
from .Unet_resnet import unet_resnet_model
from .D_LinkNet import model_DinkNet34
from .hourglass import PoseNet
from .single_model import single_unet_resnet_model

class CPN(nn.Module):
    def __init__(self, resnet, output_shape, num_class, pretrained=True):
        super(CPN, self).__init__()
        channel_settings = [2048, 1024, 512, 256]
        self.resnet = resnet
        self.global_net = globalNet(channel_settings, output_shape, num_class)
        self.refine_net = refineNet(channel_settings[-1], output_shape, num_class)

    def forward(self, x):
        res_out = self.resnet(x)
        global_fms, global_outs = self.global_net(res_out)
        refine_out = self.refine_net(global_fms)

        return global_outs, refine_out

def CPN50(out_size,num_class,pretrained=True):
    res50 = resnet50(pretrained=pretrained)
    model = CPN(res50, output_shape=out_size,num_class=num_class, pretrained=pretrained)
    return model

def CPN101(out_size,num_class,pretrained=True):
    res101 = resnet101(pretrained=pretrained)
    model = CPN(res101, output_shape=out_size,num_class=num_class, pretrained=pretrained)
    return model

def model_cpn50(image_size = (320,512),num_class = 1):
    model = CPN50(out_size=image_size, num_class=num_class, pretrained=False).cuda()
    model_state_dict = model.state_dict()
    pre_trained_state_dict = torch.load('../pre_trained_models/CPN50_384x288.pth')['state_dict']
    pre_trained_state_dict = {k.replace('module.',''):v for k,v in pre_trained_state_dict.items() if 'predict' not in k }
    model_state_dict.update(pre_trained_state_dict)
    model.load_state_dict(model_state_dict)
    return model

def model_simple_pose_res101(pretrain = False):
    model = get_pose_net(num_layers=101).cuda()
    if pretrain:
        model_state_dict = model.state_dict()
        pre_trained_dict = torch.load('../pre_trained_models/pose_resnet_101_384x288.pth')
        pre_trained_dict = {k: v for k, v in pre_trained_dict.items() if 'final' not in k}
        model_state_dict.update(pre_trained_dict)
        model.load_state_dict(model_state_dict)
    return model

def model_simple_pose_res18(num_class,pretrain=False):
    model = get_pose_net(num_layers=18,num_class=num_class).cuda()
    if pretrain:
        model_state_dict = model.state_dict()
        pre_trianed_dict = torch.load('../pre_trained_models/resnet18.pth')
        pre_trianed_dict = {k: v for k, v in pre_trianed_dict.items() if 'fc' not in k}
        model_state_dict.update(pre_trianed_dict)
        model.load_state_dict(model_state_dict)
    return model

def model_simple_pose_res34(num_class,pretrain=False):
    model = get_pose_net(num_layers=34,num_class=num_class).cuda()
    # if pretrain:
    #     model_state_dict = model.state_dict()
    #     pre_trianed_dict = torch.load('../pre_trained_models/resnet18.pth')
    #     pre_trianed_dict = {k: v for k, v in pre_trianed_dict.items() if 'fc' not in k}
    #     model_state_dict.update(pre_trianed_dict)
    #     model.load_state_dict(model_state_dict)
    return model

def model_unet():
    model = UNet(n_class=7).cuda()
    return model

def model_resnet_unet(layer,pre_train,num_class):
    model = unet_resnet_model(layers=layer,pre_train=pre_train,num_class=num_class)
    return model

def model_DLinknet34(num_class =2,pre_trained = True,):
    model = model_DinkNet34(num_class=num_class,pre_trained = pre_trained)
    return model

def model_posenet(pre_trained = True):
    model = PoseNet().cuda()
    if pre_trained:
        model_state_dict = model.state_dict()
        pretrain_state_dict = torch.load('../../pre_trained_models/hourglass_2.pt')['state_dict']
        pretrain_state_dict = {k.replace('model.module.', ''): v for k, v in pretrain_state_dict.items() if
                               ('out' not in k and 'pred' not in k)}
        model_state_dict.update(pretrain_state_dict)
        model.load_state_dict(model_state_dict)
    return model

def model_single_unet_res34(pre_trained = False):
    model = single_unet_resnet_model(layers='resnet34',pre_train=pre_trained)
    return model


if __name__ == '__main__':

    input = torch.randn(1,3,256,512).cuda()
    model = CPN50(out_size=(256,512),num_class=7,pretrained=False).cuda()
    out = model(input)
    print(out[0][0].shape,out[0][1].shape,out[0][2].shape,out[0][2].shape,out[1].shape)